app.py
from flask import Flask, render_template, redirect, url_for, flash, session, request
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate
from flask_login import LoginManager, UserMixin, login_user, logout_user, login_required, current_user
import logging
from logging.handlers import RotatingFileHandler
app = Flask(__name__)
app.config['SECRET_KEY'] = 'your_secret_key'
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///site.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db = SQLAlchemy(app)
migrate = Migrate(app, db)
login_manager = LoginManager(app)
login_manager.login_view = 'login'
# 用戶模型
class User(db.Model, UserMixin):
id = db.Column(db.Integer, primary_key=True)
username = db.Column(db.String(150), nullable=False, unique=True)
password = db.Column(db.String(150), nullable=False)
role = db.Column(db.String(50), nullable=False, default='user')
@login_manager.user_loader
def load_user(user_id):
return User.query.get(int(user_id))
@app.route('/')
@login_required
def index():
return render_template('index.html', history=current_user.history)
@app.route('/login', methods=['GET', 'POST'])
def login():
if request.method == 'POST':
username = request.form.get('username')
password = request.form.get('password')
user = User.query.filter_by(username=username).first()
if user and user.password == password:
login_user(user)
flash('Login successful!', 'success')
return redirect(url_for('index'))
else:
flash('Login failed. Check your username and password.', 'danger')
return render_template('login.html')
@app.route('/register', methods=['GET', 'POST'])
def register():
if request.method == 'POST':
username = request.form.get('username')
password = request.form.get('password')
if User.query.filter_by(username=username).first():
flash('Username already exists.', 'danger')
return redirect(url_for('register'))
new_user = User(username=username, password=password)
db.session.add(new_user)
db.session.commit()
flash('Registration successful!', 'success')
return redirect(url_for('login'))
return render_template('register.html')
@app.route('/logout')
@login_required
def logout():
logout_user()
flash('You have been logged out.', 'success')
return redirect(url_for('login'))
@app.errorhandler(404)
def not_found_error(error):
return render_template('404.html'), 404
@app.errorhandler(500)
def internal_error(error):
db.session.rollback()
return render_template('500.html'), 500
if __name__ == '__main__':
# 設置日誌記錄
if not app.debug:
file_handler = RotatingFileHandler('error.log', maxBytes=10240, backupCount=10)
file_handler.setFormatter(logging.Formatter(
'%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'
))
file_handler.setLevel(logging.INFO)
app.logger.addHandler(file_handler)
app.logger.setLevel(logging.INFO)
app.logger.info('Application startup')
app.run(debug=True)
test_app.py
import unittest
from app import app, db, User
class FlaskTestCase(unittest.TestCase):
def setUp(self):
"""在每個測試前執行,設置測試環境"""
app.config['TESTING'] = True
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///:memory:'
self.app = app.test_client()
db.create_all()
def tearDown(self):
"""在每個測試後執行,清理測試環境"""
db.session.remove()
db.drop_all()
def test_home_page(self):
"""測試首頁是否能正確加載"""
response = self.app.get('/')
self.assertEqual(response.status_code, 200)
def test_user_registration(self):
"""測試用戶註冊功能"""
response = self.app.post('/register', data=dict(
username='testuser',
password='password'
), follow_redirects=True)
self.assertEqual(response.status_code, 200)
user = User.query.filter_by(username='testuser').first()
self.assertIsNotNone(user)
def test_user_login(self):
"""測試用戶登錄功能"""
self.app.post('/register', data=dict(
username='testuser',
password='password'
), follow_redirects=True)
response = self.app.post('/login', data=dict(
username='testuser',
password='password'
), follow_redirects=True)
self.assertEqual(response.status_code, 200)
if __name__ == '__main__':
unittest.main()
index.html
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Home</title>
</head>
<body>
<h1>Welcome, {{ current_user.username }}!</h1>
<p>Your role: {{ current_user.role }}</p>
<p><a href="{{ url_for('logout') }}">Logout</a></p>
</body>
</html>
在執行 flask db migrate 和 flask db upgrade 前,你需要設置遷移環境。首先執行:
flask db init
然後執行遷移命令:
flask db migrate -m "Add role column to User"
flask db upgrade
flask run